import common.TreeNode;

/**
 * 222. Count Complete Tree Nodes 完全二叉树节点的个数
 * https://leetcode-cn.com/problems/count-complete-tree-nodes/
 */
class CountCompleteTreeNodes {

    /**
     *     方法：countNodes 计算完全二叉树的节点数量
     *
     *     Args:
     *         root: Optional[TreeNode] - 完全二叉树的根节点
     *
     *     Returns:
     *         int: 完全二叉树的节点总数
     *
     *     Time: O(logN * logN) - 其中N是节点数，每次递归调用get_depth需要O(logN)时间，递归深度为O(logN)
     *
     *     Space: O(logN) - 递归调用栈的深度
     */
    public int countNodes(TreeNode root) {
        if (root == null){
            return 0;
        }
        int leftDepth = getDepth(root.left);
        int rightDepth = getDepth(root.right);

        if( leftDepth == rightDepth){
            // 左子树是满二叉树，可以直接计算其节点数(2^leftDepth)
            // 然后递归计算右子树
            return (1 << leftDepth) + countNodes(root.right);
        } else {
            // 右子树是满二叉树(比左子树少一层)
            // 递归计算左子树
            return (1 << rightDepth) + countNodes(root.left);
        }
    }

    private int getDepth(TreeNode node){
        int depth = 0;
        while (node != null){
            depth ++;
            node = node.left;
        }
        return depth;
    }

    /**
     *      方法：countNodes2 使用二分查找和位运算计算完全二叉树的节点数量
     *           利用完全二叉树的性质，计算可能的节点范围
     *           通过二分查找和位运算高效判断节点是否存在
     *
     *      Args:
     *          root: Optional[TreeNode] - 完全二叉树的根节点
     *
     *      Returns:
     *          int: 完全二叉树的节点总数
     *
     *      Time: O(logN * logN) - 其中N是节点数，每次exists调用需要O(logN)时间，二分查找需要O(logN)次调用
     *
     *      Space: O(1) - 只使用了常数级别的额外空间
     */
    public int countNodes2(TreeNode root) {
        if (root == null){
            return 0;
        }
        int level = 0;
        TreeNode node = root;
        while(node.left != null){
            level++;
            node = node.left;
        }
        // 节点总数 N 满足 2^h ≤ N < 2^(h+1) -1
        int low = 1 << level;
        int high = (1 << (level + 1)) - 1;
        while(low < high){
            int mid = (high - low + 1) / 2 + low;
            if(exists(root, level, mid)){
                low = mid;
            } else {
                high = mid - 1;
            }
        }
        return low;
    }

    private boolean exists(TreeNode root, int level, int k){
        int bits = 1 << (level - 1);
        TreeNode node = root;
        while(node != null && bits > 0){
            if ((bits & k) == 0){
                node = node.left;
            } else {
                node = node.right;
            }
            bits >>= 1;
        }
        return node != null;
    }
}